#!/usr/bin/env python3

import fcntl
import os
import re
import socket
import sys

import configparser
import socketserver
from socketserver import ForkingMixIn, TCPServer

import nipap_whoisd

import pynipap
from pynipap import Prefix, VRF, NipapError



def format_line(key, value):
    """ Format  a single key, value pair in a whois looking format
    """
    if value is None:
        return ''

    return "%-20s %s\n" % (key + ':', value)



def format_prefix(p):
    """ Format a pynipap Prefix object into a whois looking format
    """
    if p.family == 4:
        res = format_line("inetnum", p.prefix)
    else:
        res = format_line("inet6num", p.prefix)
    res += format_line('vrf-rt', p.vrf.rt)
    res += format_line('vrf-name', p.vrf.name)
    res += format_line('descr', p.description)
    res += format_line('comment', p.comment)
    res += format_line('country', p.country)
    res += format_line('order', p.order_id)
    res += format_line('node', p.node)
    res += format_line('monitor', p.monitor)
    res += format_line('alarm-priority', p.alarm_priority)
    return res



def format_vrf(vrf):
    """ Format a pynipap VRF object into a whois looking format
    """
    res = format_line('rt', vrf.rt)
    res += format_line('vrf', vrf.name)
    res += format_line('descr', vrf.description)
    return res


# monkey patching the ForkingTCPServer for IPv6 support
class ForkingTCPServer(ForkingMixIn, TCPServer):
    def server_bind(self):
        self.socket = socket.socket(socket.AF_INET6, self.socket_type)
        self.socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, False)
        socketserver.TCPServer.server_bind(self)


class WhoisServer(socketserver.StreamRequestHandler):
    def handle(self):
        """ Called for every connection
        """
        # receive data
        query = self.rfile.readline().decode("utf-8").strip()

        try:
            query = query.decode("idna")
        except:
            pass

        # respond!
        total_hits = 0
        response = '% This is the NIPAP whois query service.\n\n'

        if query == '':
            pass
        elif re.match('^[0-9]+:[0-9]+:.*$', query):
            # 123:123:1.2.3.4 - IP address and VRF
            # respond with prefixes
            m = re.match('^([0-9]+:[0-9]+):(.*)$', query)
            vrf_query = {
                    'val1': 'vrf_rt',
                    'operator': 'equals',
                    'val2': m.group(1)
                }
            res = Prefix.smart_search(m.group(2), { }, vrf_query)
            if len(res['result']) > 0:
                total_hits += len(res['result'])
                for p in res['result']:
                    response += format_prefix(p)
                    response += '\n'
        else:
            # respond with VRFs
            res = VRF.smart_search(query, { })
            if len(res['result']) > 0:
                total_hits += len(res['result'])
                response += '% VRFs matching\n'
                for p in res['result']:
                    response += format_vrf(p)
                    response += '\n'

            # respond with prefixes
            res = Prefix.smart_search(query, { })
            if len(res['result']) > 0:
                total_hits += len(res['result'])
                for p in res['result']:
                    response += format_prefix(p)
                    response += '\n'

        if total_hits == 0:
            response += "%ERROR:101: no entries found\n"
            response += "%\n"
            response += "% No entries found in source NIPAP.\n"
            response += '%\n'
            response += '% Query the whois service by IP address, prefix, VRF RT or free text\n'
            response += '% search. It is possible to search for a prefix in a VRF by\n'
            response += '% concatenating the VRF RT and the prefix: 123:123:1.2.3.4\n'
            response += '%\n'

        response += "\n%% This query was served by the NIPAP whois server version %s\n" % nipap_whoisd.__version__
        self.request.send(response.encode("utf-8"))



if __name__ == '__main__':

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", dest="config_file", default="/etc/nipap/nipap.conf", help="read configuration from file CONFIG_FILE")
    parser.add_argument("--foreground", action="store_true", help="run in foreground")
    parser.add_argument("--listen", metavar="ADDRESS", help="listen on IPv4/6 ADDRESS")
    parser.add_argument("--port", help="port to listen to")
    parser.add_argument("--pid-file", help="write a PID file to PID_FILE")
    parser.add_argument("--no-pid-file", action="store_true", default=False, help="turn off writing PID file (overrides config file)")
    parser.add_argument("--version", action="store_true", help="display version information and exit")

    args = parser.parse_args()

    if args.version:
        print("nipap-whoisd version:", nipap_whoisd.__version__)
        sys.exit(0)

    default = {
        'pid_file': '',
        'listen': '0.0.0.0',
        'port': '43',
        'nipapd_host': 'localhost',
        'nipapd_port': '1337',
        'user': '',
        'group': ''
    }
    cfg = configparser.ConfigParser(default)
    cfg.read(args.config_file)

    # override config with command line args
    if args.listen:
        cfg.set('whoisd', 'listen', args.listen)
    if args.port:
        cfg.set('whoisd', 'port', args.port)
    if args.pid_file:
        cfg.set('whoisd', 'pid_file', args.pid_file)

    socketserver.TCPServer.allow_reuse_address = True
    try:
        server = ForkingTCPServer((cfg.get('whoisd', 'listen'),
            int(cfg.get('whoisd', 'port'))), WhoisServer)
    except socket.error as exc:
        print("Unable to bind to socket", str(exc), file=sys.stderr)
        sys.exit(1)

    # drop privileges
    if cfg.get('whoisd', 'user') != '':
        run_user = cfg.get('whoisd', 'user')
        if cfg.get('whoisd', 'group') != '':
            run_group = cfg.get('whoisd', 'group')
        else:
            run_group = cfg.get('whoisd', 'user')
        try:
            nipap_whoisd.drop_privileges(run_user, run_group)
        except Exception as e:
            print(("nipap-whoisd is configured to drop privileges and run as user '%s' and group '%s', \n"
                        "but was not started as root and can therefore not drop privileges") % (run_user, run_group), file=sys.stderr)
            sys.exit(1)
        except KeyError:
            print("Could not drop privileges to user '%s' and group '%s'" % (run_user, run_group), file=sys.stderr)
            sys.exit(1)

    # daemonize
    if not args.foreground:
        nipap_whoisd.createDaemon()

    # pid file handling
    if cfg.get('whoisd', 'pid_file') != '' and not args.no_pid_file:
        # need a+ to be able to read PID from file
        try:
            lf = open(cfg.get('whoisd', 'pid_file'), 'a+', 0)
        except IOError as exc:
            print("Unable to open PID file '%s': %s" % (
                    str(exc.filename), str(exc.strerror)), file=sys.stderr)
            sys.exit(1)
        try:
            fcntl.flock(lf, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except IOError:
            print('nipap-whoisd already running (pid: %s)' % lf.read().strip(), file=sys.stderr)
            sys.exit(1)
        print('Writing PID to file: %s' % cfg.get('whoisd', 'pid_file'), file=sys.stderr)
        lf.truncate()
        lf.write('%d\n' % os.getpid())
        lf.flush()

    pynipap.xmlrpc_uri = 'http://%(username)s:%(password)s@%(host)s:%(port)s' % {
            'username': cfg.get('whoisd', 'nipapd_username'),
            'password': cfg.get('whoisd', 'nipapd_password'),
            'host': cfg.get('whoisd', 'nipapd_host'),
            'port': cfg.get('whoisd', 'nipapd_port')
        }
    ao = pynipap.AuthOptions({'authoritative_source': 'nipap'})

    # and serve!
    try:
        server.serve_forever()
    except:
        server.shutdown()
